package com.gitee.c0dehvb.reactive.sample.r2dbc;

import com.google.common.collect.Lists;
import org.apache.shardingsphere.sql.parser.SQLParserEngine;
import org.apache.shardingsphere.sql.parser.binder.SQLStatementContextFactory;
import org.apache.shardingsphere.sql.parser.binder.metadata.schema.SchemaMetaData;
import org.apache.shardingsphere.sql.parser.binder.statement.SQLStatementContext;
import org.apache.shardingsphere.sql.parser.sql.statement.SQLStatement;
import org.apache.shardingsphere.underlying.rewrite.SQLRewriteEntry;
import org.apache.shardingsphere.underlying.rewrite.context.SQLRewriteContext;
import org.apache.shardingsphere.underlying.rewrite.engine.SQLRewriteEngine;
import org.apache.shardingsphere.underlying.rewrite.engine.SQLRewriteResult;
import org.apache.shardingsphere.underlying.rewrite.engine.SQLRouteRewriteEngine;
import org.apache.shardingsphere.underlying.route.DataNodeRouter;
import org.apache.shardingsphere.underlying.route.context.RouteContext;
import org.apache.shardingsphere.underlying.route.context.RouteUnit;

import java.util.Collection;
import java.util.List;
import java.util.Map;

/**
 * @author LiYilin
 * @since 2022/6/16 1:59 PM
 **/
public class ShardingRewriteEngine {
    private final SchemaMetaData schemaMetaData;
    private final SQLParserEngine parserEngine;
    private final DataNodeRouter dataNodeRouter;
    private final SQLRewriteEntry sqlRewriteEntry;

    public ShardingRewriteEngine(SchemaMetaData schemaMetaData, SQLParserEngine parserEngine,
                                 DataNodeRouter dataNodeRouter, SQLRewriteEntry sqlRewriteEntry) {
        this.schemaMetaData = schemaMetaData;
        this.parserEngine = parserEngine;
        this.dataNodeRouter = dataNodeRouter;
        this.sqlRewriteEntry = sqlRewriteEntry;
    }

    public String rewriteSql(String sql, List<Object> parameters) {
        SQLStatement sqlStatement = parserEngine.parse(sql, false);
        SQLStatementContext sqlStatementContext = SQLStatementContextFactory.newInstance(schemaMetaData, sql,
                                                                                         parameters,
                                                                                         sqlStatement);
        RouteContext routeContext = dataNodeRouter.route(sql, parameters, false);
        SQLRewriteContext sqlRewriteContext = sqlRewriteEntry.createSQLRewriteContext(sql, parameters,
                                                                                      sqlStatementContext,
                                                                                      routeContext);
        Collection<String> sqls = routeContext.getRouteResult().getRouteUnits().isEmpty()
                ? rewrite(sqlRewriteContext)
                : rewrite(routeContext, sqlRewriteContext);
        return sqls.stream().findFirst().orElse(null);
    }

    private Collection<String> rewrite(final SQLRewriteContext sqlRewriteContext) {
        SQLRewriteResult sqlRewriteResult = new SQLRewriteEngine().rewrite(sqlRewriteContext);
        return Lists.newArrayList(sqlRewriteResult.getSql());
    }

    private Collection<String> rewrite(final RouteContext routeContext, final SQLRewriteContext sqlRewriteContext) {
        Collection<String> result = Lists.newArrayList();
        for (Map.Entry<RouteUnit, SQLRewriteResult> entry : new SQLRouteRewriteEngine().rewrite(sqlRewriteContext
                , routeContext.getRouteResult()).entrySet()) {
            result.add(entry.getValue().getSql());
        }
        return result;
    }
}
